from pymysql import *
from scapy.all import *
from scapy.layers import http
from time import *
#from dpkt import *
from data_type import *
from function import *
#from datetime import *
import pickle
import os

from dnsDetection import *

def ExtractKeyInfo(pkt: Packet):
    time = ''
    ptype = ''      #命名只是为了避开关键字
    src = '0.0.0.0'
    dst = '0.0.0.0'
    sport = '0'
    dport = '0'
    if pkt.IsTcp():
        time = pkt.GetTcp().GetField('time')
        sport = pkt.GetTcp().GetField('sport')
        dport = pkt.GetTcp().GetField('dport')
        
        if sport == '80' or dport == '80':
            ptype = 'http'
        elif sport == '443' or dport == '443':
            ptype = 'https'
        else:
            ptype = 'tcp'
            
        if pkt.IsIpv4():
            src = pkt.GetIpv4().GetField('srcIp')
            dst = pkt.GetIpv4().GetField('dstIp')
        elif pkt.IsIpv6():
            src = pkt.GetIpv6().GetField('srcIp')
            dst = pkt.GetIpv6().GetField('dstIp')
            
    elif pkt.IsUdp():
        time = pkt.GetUdp().GetField('time')
        sport = pkt.GetUdp().GetField('sport')
        dport = pkt.GetUdp().GetField('dport')
        
        if pkt.IsDns():
            ptype = 'dns'
        else:
            ptype = 'udp'
        
        if pkt.IsIpv4():
            src = pkt.GetIpv4().GetField('srcIp')
            dst = pkt.GetIpv4().GetField('dstIp')
        elif pkt.IsIpv6():
            src = pkt.GetIpv6().GetField('srcIp')
            dst = pkt.GetIpv6().GetField('dstIp')
            
    elif pkt.IsArp():
        time = pkt.GetArp().GetField('time')
        ptype = 'arp'
        src = pkt.GetArp().GetField('psrc')
        dst = pkt.GetArp().GetField('pdst')
            
    elif pkt.IsIcmpv4():
        time = pkt.GetIcmpv4().GetField('time')
        ptype = 'icmpv4'
        
    elif pkt.IsIcmpv6():
        time = pkt.GetIcmpv6().GetField('time')
        ptype = 'icmpv6'
        
    return (time, ptype, src, dst, sport, dport)

def SaveToDatabase(keyInfo: tuple):
    try:
        conn = connect(host='127.0.0.1', port=3306, user='root', password='020903chen..', db='protocol_infomation')
    except:
        print("连接失败！")
        
    cur = conn.cursor()
    
    #tableName = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    tableName = "keyInfo"
    
    # for pkt in flowList:
    #     keyInfo = ExtractKeyInfo(pkt)
    try:
        createTable = """CREATE TABLE IF NOT EXISTS `%s` (id INT NOT NULL AUTO_INCREMENT,
                            T TIME,
                            TYPE VARCHAR(10),
                            SRC VARCHAR(40),
                            DST VARCHAR(40),
                            SPORT VARCHAR(10),
                            DPORT VARCHAR(10),
                            NUM INT(10),
                            PRIMARY KEY(id))
                            """ %tableName
        cur.execute(createTable)
        
    except Exception as e:
        print(e)
        pass

    i = cur.execute(
        'select * from `%s` where T="%s" and TYPE="%s" and SRC="%s" and DST="%s" and SPORT="%s" and DPORT="%s"' % (
            tableName, keyInfo[0], keyInfo[1], keyInfo[2], keyInfo[3], keyInfo[4], keyInfo[5]))
    # if i == 1:
    #     cur.execute(
    #         'update `%s` set NUM=NUM+1 where T="%s" and TYPE="%s" and SRC="%s" and DST="%s" and SPORT="%s" and DPORT="%s"' % (
    #             tableName, keyInfo[0], keyInfo[1], keyInfo[2], keyInfo[3], keyInfo[4], keyInfo[5]))
    # else:
    #     cur.execute(
    #         'insert into `%s`(T, TYPE, SRC, DST, SPORT, DPORT, NUM) values("%s", "%s", "%s", "%s", "%s", "%s", 1)' % (
    #             tableName, keyInfo[0], keyInfo[1], keyInfo[2], keyInfo[3], keyInfo[4], keyInfo[5]))
    cur.execute(
        'insert into `%s`(T, TYPE, SRC, DST, SPORT, DPORT, NUM) values("%s", "%s", "%s", "%s", "%s", "%s", 1)' % (
            tableName, keyInfo[0], keyInfo[1], keyInfo[2], keyInfo[3], keyInfo[4], keyInfo[5]))

    conn.commit()
    cur.close()
    conn.close()  

def PacketHandler(pkt):
    protocolInfoList = repr(pkt).split('|')
    #print(protocolInfoList)

    Ethernet = EthernetPkt()
    Ipv4 = Ipv4Pkt()
    Ipv6 = Ipv6Pkt()
    Tcp = TcpPkt()
    Udp = UdpPkt()
    Arp = ArpPkt()
    Dns = DnsPkt()
    Icmpv4 = Icmpv4Pkt()
    Icmpv6 = Icmpv6Pkt()
    HttpReq = HttpReqPkt()
    HttpRes = HttpResPkt()
    
    thisPkt = Packet(Ethernet, Ipv4, Ipv6, Tcp, Udp, Arp, Dns, Icmpv4, Icmpv6, HttpReq, HttpRes)

    # if pkt.haslayer('HTTP'):
    # pkt.show()
    GetOnePkt(protocolInfoList, thisPkt)
    
    #print(thisPkt)
    #thisPkt.PrintFlag()
    
    #forCur = open("C://Users//leslie//Desktop//res//cur//demo.pkl", 'ab') #根据实际情况替换称自己的路径
    forCur = open("sniffRes//cur//demo.pkl", 'ab')
    thisPktTrans = pickle.dumps(thisPkt)
    forCur.write(thisPktTrans)
    forCur.close()

    forCur = open("sniffRes//cur//dns.pkl", 'ab')
    pktTrans = pickle.dumps(pkt)
    forCur.write(pktTrans)
    forCur.close()
    
    #forAfter = open("C://Users//leslie//Desktop//res//after//demo.pkl", 'ab') #根据实际情况替换称自己的路径
    forAfter = open("sniffRes//after//demo.pkl", 'ab')
    thisPktTrans = pickle.dumps(thisPkt)
    forAfter.write(thisPktTrans)
    forAfter.close()
    
    keyInfo = ExtractKeyInfo(thisPkt)
    SaveToDatabase(keyInfo)

def ClearKeyInfoTable():
    try:
        conn = connect(host='127.0.0.1', port=3306, user='root', password='020903chen..', db='protocol_infomation')
    except:
        print("清空时连接失败！")
        
    cur = conn.cursor()
    tableName = "keyInfo"
    
    try:
        clearTable = """TRUNCATE TABLE `%s`""" %tableName
        cur.execute(clearTable)
    except Exception as e:
        print(e)
        pass
    
    conn.commit()
    cur.close()
    conn.close() 

def startSniff(interface, sniffCount: int):                     
    # if os.path.exists("C://Users//leslie//Desktop//res//cur//demo.pkl"): #根据实际情况替换称自己的路径
    #     os.remove("C://Users//leslie//Desktop//res//cur//demo.pkl")
    # path0="sniffRes//cur//demo.pkl"
    # path1="sniffRes//cur//dns.pkl"
    # if not os.path.exists(path0):
    #     os.makedirs(path0)
    # if not os.path.exists(path1):
    #     os.makedirs(path1)
    if os.path.exists("sniffRes//cur//demo.pkl"):
        os.remove("sniffRes//cur//demo.pkl")
    if os.path.exists("sniffRes//cur//dns.pkl"):
        os.remove("sniffRes//cur//dns.pkl")
    ClearKeyInfoTable()               #清空数据库表
    sniff(iface=interface, prn=PacketHandler, count=sniffCount)
    print("开始抓包成功")

#if __name__ == '__main__':
    #startSniff()
    #sniff(offline="submission.pcap",prn=PacketHandler,store=0)
    #print("开始抓包成功")
    #dnsDetection()
    #print(mes)
    #pktCount = 0
    ##ClearKeyInfoTable()
    ##flow = []
    #package = sniff(filter="port 53", iface="Intel(R) Wireless-AC 9560 160MHz", prn=PacketHandler, count=10)

    #package = sniff(prn=PacketHandler, count=0)
    # packet10 = flow[9]
    # if packet10.IsIpv4():
    #     print("here!!!")
    #     packet10.ipv4.PrintMyInfo()
    # print(flow)
    #wrpcap("C://Users//leslie//Desktop//framework//demo.pcap", package)
    ###package = sniff(prn=PacketHandler, count=10)
    ##print(len(flow))
    ##wrpcap("C://Users//leslie//Desktop//res//demo.pcap", package)
    # SaveToDatabase(flow)
    # for pkt in flow:
    #     pkt.PrintFlag()
    # print(len(flow))
    # packet10 = flow[9]
    # if packet10.IsIpv4():
    #     print("here!!!")
    #     packet10.ipv4.PrintMyInfo()
